from functools import lru_cache
from web3 import Web3
from coin.blockchain.web3.eth.constants import (
  ETH_TOKEN_TO_ADDRESS, ERC20_ABI)


class EthNativePrivateClient:
  def __init__(self, *, node_info, api_key=None, timeout=60):
    self._node_info = node_info
    self._web3 = Web3(Web3.HTTPProvider(
        self._node_info,
        request_kwargs={"timeout": timeout}))

  def get_eth_balance(self, address, end_height=None):
    if end_height is not None:
      return self._web3.eth.get_balance(address, end_height)
    else:
      return self._web3.eth.get_balance(address)

  @lru_cache(maxsize=20)
  def _load_contract_erc20(self, address):
    return self._web3.eth.contract(address=address, abi=ERC20_ABI)

  def get_token_balance(self, address, token_address, end_height=None):
    if ETH_TOKEN_TO_ADDRESS['ETH'] == token_address:
      return self.get_eth_balance(
          Web3.to_checksum_address(address), end_height)
    else:
      kwargs = {"block_identifier": end_height} if end_height else {}
      contract_erc20 = self._load_contract_erc20(token_address)
      return contract_erc20.functions.balanceOf(
          Web3.to_checksum_address(address)).call(**kwargs)

  def get_latest_block_number_timestamp(self):
    block = self._web3.eth.get_block('latest', False)
    block_num = block['number']
    block_time = block['timestamp'] * 10 ** 9
    return block_num, block_time
  
  def get_block_timestamp(self, block_number):
    return self._web3.eth.get_block(block_number, False)['timestamp'] * 10 ** 9
